package com.example.springweb.rest;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import com.alibaba.fastjson.JSONObject;

import cn.hutool.core.util.StrUtil;
import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.tags.Tag;
import reactor.core.publisher.Flux;

@Tag(name = "智能AI大模型接口")
@RestController
@RequestMapping(value = "/ai")
public class ChatController {
	
	// 存储每个用户的对话上下文 Prompt 对象
	private final Map<String, List<UserMessage>> userMap = new ConcurrentHashMap<>();
	
	@Autowired
	private OllamaChatModel ollamaChatModel;
	
	@Schema(title = "普通对话聊天")
	@GetMapping(value = "/chat")
	public String chat(@RequestParam("message") String message,
			@RequestParam("userId") String userId) {
		if(StrUtil.isEmptyIfStr(userId)) {
			return "您什么都没有说...";
		}
		if(message.length() > 1000) {
			return "不能超过1000个文字(含标点符号)";
		}
		// 获取用户的消息列表，如果不存在则创建一个新的
        List<UserMessage> messages = userMap.computeIfAbsent(userId, k -> new ArrayList<>());
        
        // 添加新的用户消息
        UserMessage userMessage = new UserMessage(message);
        messages.add(userMessage);
        
        // 根据消息列表构建 Prompt 对象
        Prompt prompt = new Prompt(userMessage);
		
		ChatResponse chatResponse = ollamaChatModel.call(prompt);
		return chatResponse.toString();
	}

	@Schema(title = "流式响应对话聊天")
	@GetMapping(value = "/streamChat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
	public Flux<ChatResponse> streamChat(@RequestParam("message") String message,
			@RequestParam("userId") String userId) {
		// 获取用户的消息列表，如果不存在则创建一个新的
        List<UserMessage> messages = userMap.computeIfAbsent(userId, k -> new ArrayList<>());
        
        // 添加新的用户消息
        UserMessage userMessage = new UserMessage(message);
        messages.add(userMessage);
        
        // 根据消息列表构建 Prompt 对象
        Prompt prompt = new Prompt(userMessage);
        List<Message> messageList = prompt.getInstructions();
        messageList.forEach(m -> System.out.println(userId + "，问：" + m.getText()));
        return ollamaChatModel.stream(prompt);
	}
	
	@Schema(title = "流式响应对话聊天")
	@GetMapping(value = "/test", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
	public Flux<String> test(){
		Flux<String> messages = Flux.just(" Message 1", " Message 2 ", " Message 3", " Message 4");
		messages.toStream().forEach(System.out::println);
		return messages.delayElements(Duration.ofSeconds(1));
	}
}
